package com.ivyft.jetty.yarn.protocol;

import com.ivyft.jetty.yarn.JettyAMRMClient;
import com.ivyft.jetty.yarn.JettyAppMaster;
import com.ivyft.jetty.yarn.JettyConfiguration;
import com.ivyft.jetty.yarn.Shutdown;
import org.apache.avro.AvroRemoteException;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * <pre>
 *
 * Created by IntelliJ IDEA.
 * User: zhenqin
 * Date: 15/11/27
 * Time: 20:59
 * To change this template use File | Settings | File Templates.
 *
 * </pre>
 *
 * @author zhenqin
 */
public class JettyYarnMasterProtocol implements JettyYarnProtocol {


    protected final JettyConfiguration conf;


    protected final JettyAMRMClient amrmClient;



    protected JettyAppMaster appMaster;


    private static Logger LOG = LoggerFactory.getLogger(JettyYarnMasterProtocol.class);



    public JettyYarnMasterProtocol(JettyConfiguration conf, JettyAMRMClient client) {
        this.conf = conf;
        this.amrmClient = client;
    }

    @Override
    public Void shutdown() throws AvroRemoteException {
        LOG.info("remote shutdown...");
        if(appMaster != null) {
            LOG.info("remote shutdown message...");
            appMaster.add(new Shutdown());
        }
        return null;
    }

    @Override
    public List<NodeContainer> listInstance() throws AvroRemoteException {
        return this.amrmClient.listJettyNodes();
    }

    @Override
    public Void addInstance(int memory, int cores,
                            CharSequence jettyZip,
                            List<CharSequence> warZips,
            java.lang.CharSequence jettyXml) throws AvroRemoteException {
        if(amrmClient.containerIsEmpty()) {
            this.appMaster.newContainer(memory, cores);
        }
        Set<String> wars = new HashSet<String>(this.appMaster.getWars());
        for (CharSequence warZip : warZips) {
            wars.add(warZip.toString());
        }
        jettyZip = jettyZip == null ? "" : jettyZip;
        jettyXml = jettyXml == null ? "./conf/jetty.xml" : jettyXml;
        this.amrmClient.startInstance(jettyZip.toString(), new ArrayList<String>(wars), jettyXml.toString());
        return null;
    }

    @Override
    public Void register(CharSequence containerid, CharSequence nodeId, int port) throws AvroRemoteException {
        return null;
    }

    @Override
    public Void stopInstance(CharSequence id, IdType idType) throws AvroRemoteException {
        if(idType == IdType.CONTAINER_ID) {
            try {
                this.amrmClient.stopInstance(id.toString());
                return null;
            } catch (Exception e) {
                throw new AvroRemoteException(e);
            }
        } else {
            try {
                List<NodeContainer> nodeContainers = listInstance();
                for (NodeContainer nodeContainer : nodeContainers) {
                    if (StringUtils.equals(String.valueOf(id), String.valueOf(nodeContainer.getNodeHost()))) {
                        this.amrmClient.stopInstance(nodeContainer.getContainerId().toString());
                    }
                }
            } catch (Exception e) {
                throw new AvroRemoteException(e);
            }

            return null;
        }
    }

    @Override
    public Void stopAll() throws AvroRemoteException {
        List<NodeContainer> nodeContainers = listInstance();
        for (NodeContainer container : nodeContainers) {
            try {
                this.amrmClient.stopInstance(container);
            } catch (Exception e) {
                throw new AvroRemoteException(e);
            }
        }
        return null;
    }

    @Override
    public Void close() throws AvroRemoteException {
        return null;
    }


    public JettyConfiguration getConf() {
        return conf;
    }

    public JettyAMRMClient getAmrmClient() {
        return amrmClient;
    }

    public JettyAppMaster getAppMaster() {
        return appMaster;
    }

    public void setAppMaster(JettyAppMaster appMaster) {
        this.appMaster = appMaster;
    }
}
